'''
Author: zhanwei xu
Date: 2023-06-10 10:41:13
LastEditors: zhanwei xu
LastEditTime: 2024-02-02 11:05:01
Description: 

Copyright (c) 2023 by zhanwei xu, Tsinghua University, All Rights Reserved. 
'''
import pymongo as mg
import requests
import httpx
import os
from typing import Any, Dict, List, Optional
from pydantic import BaseModel
from fastapi import APIRouter, Request, HTTPException, UploadFile, File
from fastapi.responses import JSONResponse, StreamingResponse
from datetime import datetime, timedelta
from core import setting, mgClient
from uuid import uuid1
from sys import getsizeof
import asyncio
import functools
from .utils import *
import random
import shutil
from concurrent.futures import ThreadPoolExecutor
db = mgClient['chat']
collection = db.get_collection('file')

current_timestamp = int(datetime.now().timestamp()*1000)
collection.insert_one({
    "ip": 0,
    "timestamp": {"$lt": current_timestamp}
})

collection_file = db.get_collection('file_db')
router = APIRouter(
    prefix="/share",
    tags=["share"]
)


def check_ip_request_count(ip: str) -> bool:
    """
    It checks the traffic of the given IP address.
    
    :param ip: The IP address of the host to check
    """
    now = datetime.now()
    traffic = collection.find_one({"ip": ip})  # 查找指定IP地址的流量信息
    if not traffic:
        traffic = {"ip": ip, "data": [{"timestamp": now, "count": 1}]}
        collection.insert_one(traffic)  # 插入新的流量信息
        return False
    else:
        recent_count = sum([t["count"] for t in traffic["data"]])
        # 删除超过24小时的旧流量数据
        traffic["data"] = [t for t in traffic["data"] if now - t["timestamp"] <= timedelta(hours=24)]
        if recent_count >= 30:  # 设置流量限制为10分钟内最多60个请求
            return True
        else:
            traffic["data"].append({"timestamp": now, "count": 1})
            collection.update_one({"ip": ip}, {"$set": traffic})  # 更新流量信息
            return False

def get_clients():
    f = open('p_b_value.txt', 'r')
    lines = f.readlines()
    f.close()
    line_index = random.randint(0, len(lines) - 1)
    token = lines[line_index].strip()
    return token


def get_gemini_keys():
    f = open('gemini_keys.txt', 'r')
    lines = f.readlines()
    f.close()
    line_index = random.randint(0, len(lines) - 1)
    token = lines[line_index].strip()
    return token

def get_user_key():
    f = open('user_key.txt', 'r')
    lines = f.readlines()
    f.close()
    line_index = random.randint(0, len(lines) - 1)
    token = lines[line_index].strip()
    return token

async def async_post(url, data=None, json = None, headers=None, timeout=10):
    loop = asyncio.get_running_loop()
    request = functools.partial(requests.post, url, data=data,json = json)
    try:
        with ThreadPoolExecutor() as executor:
            response = await loop.run_in_executor(
                executor, request
            )
    except:
        # 返回一个空的response，状态码为404
        response = requests.Response()
        response.status_code = 500
    return response



@router.post("/message")
async def message(request: Request) -> Optional[str]:
    client_token = get_clients()
    json_data = await request.json()
    assert_id_list = json_data["assets"]
    filepath_list = []
    assert_infos = get_multiple_assets(assert_id_list)
    assert_infos = assert_infos["list"]
    for assert_info in assert_infos:
        name = assert_info["name"]
        download_url = assert_info["url"]
        assert_id = assert_info["id"]
        file_extension = os.path.splitext(name)[1]
        unique_filename = f"{assert_id}{file_extension}"
        file_path = os.path.join("uploads", unique_filename)
        filepath_list.append(file_path)
        # 检查文件是否已存在
        if os.path.exists(file_path):
            print(f"文件已存在，跳过下载: {file_path}")
            continue
        
        
        loop = asyncio.get_running_loop()
        request = functools.partial(requests.get, download_url)
        try:
            with ThreadPoolExecutor() as executor:
                response = await loop.run_in_executor(
                    executor, request
                )
        except:
            print(response.text)
            response = requests.Response()
            response.status_code = 500
        # # 下载文件
        # response = requests.get(download_url)
        # 检查请求是否成功
        if response.status_code == 200:
            # 将内容写入文件
            with open(file_path, 'wb') as file:
                file.write(response.content)
            print(f"文件已下载并保存为: {file_path}")
        else:
            print("文件下载失败，状态码:", response.status_code)
    try:
        model = json_data["model"]
    except:
        model = "poe"
    content =  ""
    for filepath in filepath_list:
        content += await get_file_content(filepath)
    new_filename = f"{uuid1()}.txt"
    newfile_path = os.path.join("uploads", new_filename)
    with open(newfile_path, 'w', encoding='utf-8') as newfile:
        newfile.write(content)
    message = json_data["messages"]
    if model =="poe":
        question = message[-1]["content"]
        return StreamingResponse(generate_text_poe(client_token,newfile_path, question), media_type="text/plain")
    elif model =="gpt4-128k":
        user_key = get_user_key()
        
        user_key = json_data["key"]
        question = message[-1]["content"]
        response  = await async_request_4(question,content,message[:-1], user_key)
        # return StreamingResponse("功能还在更新中，敬请期待")
        return StreamingResponse(chattext_chatgpt(response), media_type="text/plain")
    else:
        raise HTTPException(status_code=422, detail="请传入一个正确的模型参数")

client_token = get_gemini_keys()
file_path = "在问使用指南.docx"
gemini_df = get_chunk_from_document(client_token, file_path, chunk_size=1000, chunk_overlap=20)

@router.post("/ai_customer_service")
async def message(request: Request) -> Optional[str]:
    json_data = await request.json()
    question = json_data["question"]
    return StreamingResponse(generate_text_gemini(client_token, question, gemini_df), media_type="text/plain")



